import datasets.datasets_ws as datasets_ws
import cv2
import torch

class DatasetSP(datasets_ws.BaseDataset):
    def __init__(self, context, datasets_folder="datasets", dataset_name="pitts30k", split="train"):
        super().__init__(context.args, datasets_folder, dataset_name, split)
        self._resize_flag = False
        if "resize" in context.config["data"]["preprocessing"].keys():
            self._resize_flag = True
            self.img_size = context.config["data"]["preprocessing"]["resize"]
    
    def __getitem__(self, index):
        imgpath = self.images_paths[index]
        input_image = self.read_image(imgpath)
        input_image = input_image.astype('float32')
        return input_image,index
    
    def read_image(self, impath):
        """ Read image as grayscale and resize to img_size.
        Inputs
          impath: Path to input image.
          img_size: (W, H) tuple specifying resize size.
        Returns
          grayim: float32 numpy array sized H x W with values in range [0, 1].
        """
        grayim = cv2.imread(impath, 0)
        if grayim is None:
            raise Exception('Error reading image %s' % impath)
        # Image is resized via opencv.
       
        if self._resize_flag == True:
          interp = cv2.INTER_AREA
          img_size = self.img_size
          # size(width,height)
          grayim = cv2.resize(grayim, (img_size[0], img_size[1]), interpolation=interp)
        grayim = (grayim.astype('float32') / 255.)
        return grayim